#!/usr/bin/env python
"""
数据库初始化模块
提供数据库表结构创建和管理功能
"""
import os
import sys
import argparse
from typing import Dict, Any

# 添加项目根目录到Python路径
sys.path.append(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))))

# 导入必要模块
from backend.logger_setup import setup_logging, get_logger

# 延迟导入数据库管理器以避免循环导入
from importlib import import_module

# 初始化日志
setup_logging()
logger = get_logger('database.init')

def get_db_manager():
    """获取数据库管理器实例，使用新的database模块"""
    try:
        # 直接导入新的database模块中的db_manager
        from backend.database.db import db_manager
        return db_manager
    except ImportError as e:
        logger.error(f"无法导入数据库管理器: {str(e)}")
        # 导入失败时创建一个简单的内存模式数据库管理器
        from backend.database.db import DatabaseManager
        return DatabaseManager()

# 颜色定义
class Colors:
    GREEN = '\033[92m'
    YELLOW = '\033[93m'
    RED = '\033[91m'
    RESET = '\033[0m'


def init_db(force: bool = False) -> bool:
    """初始化数据库，创建所有必要的表和索引
    
    Args:
        force: 是否强制重新初始化（删除现有表）
        
    Returns:
        bool: 初始化是否成功
    """
    try:
        db_manager = get_db_manager()
        logger.info(f"{Colors.YELLOW}开始数据库初始化...{Colors.RESET}")
        
        if force:
            logger.warning(f"{Colors.YELLOW}强制模式：将删除现有表并重新创建{Colors.RESET}")
            if not db_manager.connection:
                db_manager._connect()
            
            # 检查是否在内存模式
            if hasattr(db_manager, 'memory_documents'):
                logger.info(f"{Colors.GREEN}当前处于内存模式，弱重置所有内存表{Colors.RESET}")
                db_manager._enable_memory_mode()
            else:
                # 执行表删除操作
                cursor = db_manager.connection.cursor()
                try:
                    # 先删除索引
                    logger.info("删除现有索引...")
                    cursor.execute("DROP INDEX IF EXISTS idx_document_vectors_embedding")
                    
                    # 删除表（按依赖顺序）
                    logger.info("删除现有表...")
                    tables = ['document_vectors', 'documents', 'messages', 'session_model_map']
                    for table in tables:
                        cursor.execute(f"DROP TABLE IF EXISTS {table} CASCADE")
                    db_manager.connection.commit()
                finally:
                    cursor.close()
        
        # 初始化数据库
        db_manager._initialize_tables()
        
        # 检查表是否成功创建
        if not hasattr(db_manager, 'memory_documents') and db_manager.connection:
            cursor = db_manager.connection.cursor()
            try:
                # 列出所有表（SQLite方式）
                cursor.execute("SELECT name FROM sqlite_master WHERE type='table' ORDER BY name")
                tables = cursor.fetchall()
                
                logger.info(f"{Colors.GREEN}数据库初始化成功！已创建的表：{Colors.RESET}")
                for table in tables:
                    logger.info(f"  - {table['name']}")
                logger.info(f"{Colors.RESET}")
            finally:
                cursor.close()
        else:
            logger.info(f"{Colors.GREEN}内存模式数据库初始化成功！{Colors.RESET}")
        
        return True
    except Exception as e:
        logger.error(f"{Colors.RED}数据库初始化失败: {str(e)}{Colors.RESET}")
        # 确保回滚事务
        if hasattr(db_manager, 'connection') and db_manager.connection:
            db_manager.connection.rollback()
        return False


def drop_db() -> bool:
    """删除所有数据库表
    
    Returns:
        bool: 操作是否成功
    """
    try:
        db_manager = get_db_manager()
        logger.warning(f"{Colors.RED}警告：即将删除所有数据库表！{Colors.RESET}")
        
        if not db_manager.connection:
            db_manager._connect()
        
        # 检查是否在内存模式
        if hasattr(db_manager, 'memory_documents'):
            logger.info(f"{Colors.GREEN}当前处于内存模式，重置所有内存表{Colors.RESET}")
            db_manager._enable_memory_mode()
        else:
            cursor = db_manager.connection.cursor()
            try:
                # 删除索引
                cursor.execute("DROP INDEX IF EXISTS idx_document_vectors_embedding")
                
                # 删除表（按依赖顺序）
                tables = ['document_vectors', 'documents', 'messages', 'session_model_map']
                for table in tables:
                    logger.info(f"删除表{table}...")
                    cursor.execute(f"DROP TABLE IF EXISTS {table} CASCADE")
                db_manager.connection.commit()
                
                logger.info(f"{Colors.GREEN}所有数据库表已成功删除！{Colors.RESET}")
            finally:
                cursor.close()
        
        return True
    except Exception as e:
        logger.error(f"{Colors.RED}删除数据库表失败: {str(e)}{Colors.RESET}")
        if hasattr(db_manager, 'connection') and db_manager.connection:
            db_manager.connection.rollback()
        return False


def reset_db() -> bool:
    """重置数据库（删除所有表并重新创建）
    
    Returns:
        bool: 操作是否成功
    """
    logger.info(f"{Colors.YELLOW}开始数据库重置...{Colors.RESET}")
    
    # 获取数据库管理器实例
    db_manager = get_db_manager()
    
    # 先删除所有表
    if not drop_db():
        return False
    
    # 然后重新初始化
    return init_db()


def main():
    """命令行入口点"""
    parser = argparse.ArgumentParser(description='数据库初始化和管理工具')
    
    # 子命令解析器
    subparsers = parser.add_subparsers(dest='command', help='可用命令')
    
    # init命令
    init_parser = subparsers.add_parser('init', help='初始化数据库')
    init_parser.add_argument('-f', '--force', action='store_true', help='强制重新初始化（删除现有表）')
    
    # drop命令
    subparsers.add_parser('drop', help='删除所有数据库表')
    
    # reset命令
    subparsers.add_parser('reset', help='重置数据库（删除所有表并重新创建）')
    
    # 如果没有提供命令，显示帮助
    if len(sys.argv) == 1:
        parser.print_help()
        sys.exit(1)
    
    # 解析参数
    args = parser.parse_args()
    
    # 执行相应命令
    if args.command == 'init':
        success = init_db(args.force)
    elif args.command == 'drop':
        success = drop_db()
    elif args.command == 'reset':
        success = reset_db()
    else:
        parser.print_help()
        sys.exit(1)
    
    # 根据执行结果设置退出码
    sys.exit(0 if success else 1)


if __name__ == '__main__':
    main()
